import pytesseract
from PIL import Image
import cv2
import numpy as np
import re
import os
from collections import Counter

from openai import OpenAI

class UniversalChineseOCR:
    def __init__(self,api_key,base_url):
        self.client = OpenAI(api_key=api_key, base_url=base_url)
        # 通用中文OCR配置，不使用字符白名单
        self.configs = [
            r'--oem 1 --psm 6 -c preserve_interword_spaces=1 -l chi_sim+eng',  # 保留空格
            r'--oem 1 --psm 7 -c preserve_interword_spaces=1 -l chi_sim+eng',  # 单行文本
            r'--oem 3 --psm 6 -c tessedit_char_whitelist=ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789()（）[]【】{}《》\u4e00-\u9fff -l chi_sim+eng',
            # 字符白名单
            r'--oem 1 --psm 4 -c preserve_interword_spaces=1 -l chi_sim+eng',  # 多列文本
        ]

    def preprocess_image_universal(self, image_path):
        """改进的图像预处理，增强括号识别"""
        img = cv2.imread(image_path)
        if img is None:
            raise ValueError(f"无法读取图像文件: {image_path}")

        # 1. 自动调整图像大小 - 确保足够大的分辨率来识别细节
        height, width = img.shape[:2]
        if max(height, width) > 2000:  # 如果图像太大，缩小
            scale = 1500 / max(height, width)
            new_width = int(width * scale)
            new_height = int(height * scale)
            img = cv2.resize(img, (new_width, new_height), interpolation=cv2.INTER_AREA)
        elif max(height, width) < 1000:  # 如果图像太小，放大到更大尺寸
            scale = 1800 / max(height, width)
            new_width = int(width * scale)
            new_height = int(height * scale)
            img = cv2.resize(img, (new_width, new_height), interpolation=cv2.INTER_CUBIC)

        # 2. 多种预处理方法
        processed_images = []

        # 方法1: 标准灰度 + Otsu二值化 + 形态学操作
        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        _, thresh1 = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
        # 添加形态学操作来清理噪声，但保留细小结构
        kernel = np.ones((1, 1), np.uint8)
        thresh1 = cv2.morphologyEx(thresh1, cv2.MORPH_CLOSE, kernel)
        processed_images.append(('otsu_morph', thresh1))

        # 方法2: 自适应阈值 - 更适合处理括号等细小结构
        thresh2 = cv2.adaptiveThreshold(gray, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
                                        cv2.THRESH_BINARY, 11, 2)
        processed_images.append(('adaptive_tuned', thresh2))

        # 方法3: 对比度增强 + 锐化 - 增强括号边缘
        clahe = cv2.createCLAHE(clipLimit=3.0, tileGridSize=(8, 8))
        enhanced = clahe.apply(gray)

        # 锐化操作增强边缘细节
        kernel_sharpen = np.array([[-1, -1, -1], [-1, 9, -1], [-1, -1, -1]])
        sharpened = cv2.filter2D(enhanced, -1, kernel_sharpen)

        _, thresh3 = cv2.threshold(sharpened, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
        processed_images.append(('enhanced_sharpen', thresh3))

        # 方法4: 去噪 + 直方图均衡化 - 保持细节
        denoised = cv2.fastNlMeansDenoising(gray, None, 10, 7, 21)
        equalized = cv2.equalizeHist(denoised)
        _, thresh4 = cv2.threshold(equalized, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
        processed_images.append(('equalized_denoised', thresh4))

        # 方法5: 专门针对细小字符的处理
        blurred = cv2.GaussianBlur(gray, (3, 3), 0)
        _, thresh5 = cv2.threshold(blurred, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
        processed_images.append(('fine_detail', thresh5))

        return processed_images

    def extract_text_multiple_methods(self, image_path):
        """使用多种方法提取文本"""
        processed_images = self.preprocess_image_universal(image_path)

        all_results = []

        for method_name, img in processed_images:
            for config in self.configs:
                try:
                    # 转换为PIL图像进行OCR
                    pil_img = Image.fromarray(img)
                    text = pytesseract.image_to_string(pil_img, config=config)

                    if text.strip():
                        all_results.append({
                            'method': f"{method_name}_{config.split('psm ')[1][0]}",
                            'text': text,
                            'char_count': len(text.strip()),
                            'chinese_char_count': len(re.findall(r'[\u4e00-\u9fff]', text))
                        })

                except Exception as e:
                    continue

        return all_results

    def select_best_result(self, results):
        """选择最佳的识别结果"""
        if not results:
            return None

        # 根据中文字符数量和总字符数量评分
        scored_results = []
        for result in results:
            score = result['chinese_char_count'] * 2 + result['char_count']
            # 如果包含常见关键词，加分
            keywords = ['根据', '研究', '实现', '无法', '选项', '答案']
            keyword_bonus = sum(10 for keyword in keywords if keyword in result['text'])
            score += keyword_bonus

            scored_results.append((score, result))

        # 选择分数最高的结果
        scored_results.sort(key=lambda x: x[0], reverse=True)
        return scored_results[0][1] if scored_results else None

    def extract_question_blocks(self, text):
        """改进的题目块提取"""
        lines = [line.strip() for line in text.split('\n') if line.strip()]

        # 寻找可能的题目开始行
        question_starts = []
        for i, line in enumerate(lines):
            # 题目通常包含这些特征
            if (len(line) > 5 and
                    any(char in line for char in ['根据', '研究', '问题', '题目', '？', '?']) or
                    re.search(r'^[A-D][\.\、\．:]', line) or  # 包含选项标识符
                    re.search(r'[0-9]{1,3}/[0-9]{1,3}', line)):  # 包含页码
                question_starts.append(i)

        # 提取题目块
        question_blocks = []
        for start_idx in question_starts[:5]:  # 增加检查的数量
            block_lines = []
            option_count = 0

            for i in range(start_idx, min(start_idx + 20, len(lines))):  # 增加行数范围
                line = lines[i]
                block_lines.append(line)

                # 统计选项数量 - 改进正则表达式
                if re.match(r'^[A-D][\.\、\．:\s]', line.strip()):
                    option_count += 1

                # 如果已经收集了4个选项或者遇到答案信息，停止
                if option_count >= 4 or any(keyword in line for keyword in ['正确答案', '答案', '解析']):
                    # 继续读取几行以防答案在同一区域
                    for j in range(i + 1, min(i + 4, len(lines))):
                        next_line = lines[j]
                        if any(keyword in next_line for keyword in ['正确答案', '答案', '解析']):
                            block_lines.append(next_line)
                        elif re.match(r'^[A-D][\.\、\．:\s]', next_line.strip()):
                            block_lines.append(next_line)
                    break

            # 更宽松的有效题目块判断条件
            if len(block_lines) >= 2 and (option_count >= 2 or
                                          any('答案' in line for line in block_lines)):
                question_blocks.append({
                    'start_line': start_idx,
                    'lines': block_lines,
                    'option_count': option_count,
                    'text': '\n'.join(block_lines)
                })

        return question_blocks

    def analyze_text_structure(self, text):
        """分析文本结构"""
        lines = [line.strip() for line in text.split('\n') if line.strip()]

        structure = {
            'header': [],
            'content_blocks': [],
            'questions': [],
            'options_blocks': [],
            'footer': []
        }

        current_block = []
        current_type = 'content'

        for i, line in enumerate(lines):
            # 检测各部分类型
            if i == 0 and (re.search(r'\d+/\d+', line) or '案例' in line or '题目' in line):
                structure['header'].append(line)
                continue

            # 检测选项块
            if re.match(r'^[A-D][\.\、\．\s]', line):
                if current_block and current_type == 'content':
                    structure['content_blocks'].append('\n'.join(current_block))
                    current_block = []
                current_type = 'options'
                structure['options_blocks'].append(line)
                continue

            # 检测问题行
            if '根据' in line and '研究' in line and len(line) > 10:
                if current_block:
                    structure['content_blocks'].append('\n'.join(current_block))
                current_block = [line]
                current_type = 'question'
                continue

            # 检测答案信息
            if any(keyword in line for keyword in ['正确答案', '我的答案', '参考解析']):
                structure['footer'].append(line)
                continue

            # 添加到当前块
            current_block.append(line)

        # 添加最后一个块
        if current_block:
            if current_type == 'question':
                structure['questions'].append('\n'.join(current_block))
            else:
                structure['content_blocks'].append('\n'.join(current_block))

        return structure

    def clean_with_openai(self, ocr_text):
        """使用OpenAI清理和格式化"""
        prompt = f"""
            请分析以下从图片中识别出的题目内容。图片识别可能存在错误，请修正识别错误并提取以下信息：

            1. 题目描述（理论背景等）
            2. 题目本身
            3. 所有选项（A、B、C、D等）
            4. 正确答案

            要求：
            - 修正OCR识别中的错误字符

            识别内容:
            {ocr_text}

            请按以下格式输出，不关心分析修改的内容直接按我要求的输出（题目输出完了空一个空格紧跟答案不要换行）：
            下面是输出结果:
            【题目描述】
            ...
            【题目】
            ...
            【选项】
            A. ...
            B. ...
            C. ...
            D. ...

            【答案】
            ...
            """

        response = self.client.chat.completions.create(
            model="deepseek-ai/DeepSeek-R1-0528-Qwen3-8B",
            messages=[{'role': 'user', 'content': prompt}],
            stream=True
        )

        print("OpenAI处理中...")
        full_response = ""

        for chunk in response:
            if not chunk.choices:
                continue
            if chunk.choices[0].delta.content:
                content = chunk.choices[0].delta.content
                print(content, end="", flush=True)
                full_response += content
            if hasattr(chunk.choices[0].delta, 'reasoning_content') and chunk.choices[0].delta.reasoning_content:
                reasoning = chunk.choices[0].delta.reasoning_content
                print(reasoning, end="", flush=True)
                full_response += reasoning

        return full_response

    def extract_final_content(self, openai_response):
        """提取'下面是输出结果:'之后的内容"""
        separator = "下面是输出结果:"

        if separator in openai_response:
            # 找到分隔符位置
            separator_index = openai_response.find(separator)
            # 提取分隔符之后的内容
            final_content = openai_response[separator_index + len(separator):].strip()
            return final_content
        else:
            # 如果没有找到分隔符，返回原始内容
            print("警告：未找到'下面是输出结果:'分隔符，返回完整内容")
            return openai_response

    def save_final_content(self, final_content, image_path, output_dir="final_results"):
        """保存最终内容到文件"""
        if not os.path.exists(output_dir):
            os.makedirs(output_dir)

        base_name = os.path.splitext(os.path.basename(image_path))[0]
        output_file = os.path.join(output_dir, f"{base_name}_final.txt")

        with open(output_file, 'w', encoding='utf-8') as f:
            f.write(final_content)

        print(f"\n✅ 最终内容已保存到: {output_file}")
        return output_file

    def process_single_image(self, image_path):
        """处理单张图片"""
        print(f"处理图片: {os.path.basename(image_path)}")

        # 提取文本
        ocr_text = self.extract_text_multiple_methods(image_path)

        print(f"📝 OCR原始内容: {ocr_text}")

        # 2. OpenAI处理
        openai_response = self.clean_with_openai(ocr_text)

        # 3. 提取最终内容
        print("\n🎯 提取最终内容...")
        final_content = self.extract_final_content(openai_response)

        # 4. 保存最终内容
        saved_file = self.save_final_content(final_content, image_path)

        return {
            'original_ocr': ocr_text,
            'openai_response': openai_response,
            'final_content': final_content,
            'saved_file': saved_file
        }

    def batch_process_images(self, image_folder):
        """批量处理图片文件夹"""
        image_extensions = ['.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.tif']

        # 获取所有图片文件
        image_files = []
        for file in os.listdir(image_folder):
            if any(file.lower().endswith(ext) for ext in image_extensions):
                image_files.append(os.path.join(image_folder, file))

        print(f"找到 {len(image_files)} 张图片")

        results = []
        for image_path in image_files:
            result = self.process_single_image(image_path)
            if result:
                results.append(result)
            print("-" * 50)

        return results



def main():
    # 配置API
    API_KEY = "sk-msxhrypvitaetkpthwdwvmmoekiabedcjxlbsznqzmjigpnq"  # 替换为您的API密钥
    BASE_URL = "https://api.siliconflow.cn/v1"  # 替换为您的API地址

    ocr = UniversalChineseOCR(API_KEY, BASE_URL)

    # 选择处理模式
    print("请选择处理模式:")
    print("1. 处理单张图片")
    print("2. 批量处理图片文件夹")

    choice = input("请输入选择 (1 或 2): ").strip()

    if choice == "1":
        image_path = input("请输入图片路径: ").strip()
        if not os.path.exists(image_path):
            print("文件不存在!")
            return

        result = ocr.process_single_image(image_path)
        if result:
            print(f"\n🎉 处理完成!")
            print(f"📁 最终内容保存位置: {result['saved_file']}")


    elif choice == "2":

        folder_path = input("请输入图片文件夹路径: ").strip()

        if not os.path.exists(folder_path):
            print("❌ 文件夹不存在!")

            return

        # 获取所有图片文件

        image_extensions = ['.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.tif']

        image_files = []

        for file in os.listdir(folder_path):

            if any(file.lower().endswith(ext) for ext in image_extensions):
                image_files.append(os.path.join(folder_path, file))

        print(f"📂 找到 {len(image_files)} 张图片")

        results = []

        for i, image_path in enumerate(image_files, 1):

            print(f"\n{'=' * 60}")

            print(f"🔄 处理第 {i}/{len(image_files)} 张: {os.path.basename(image_path)}")

            print('=' * 60)

            result = ocr.process_single_image(image_path)

            if result:
                results.append(result)

            # 如果不是最后一张，等待用户确认继续

            if i < len(image_files):
                input("\n⏳ 按Enter键继续处理下一张图片...")

        print(f"\n🎊 批量处理完成! 成功处理 {len(results)}/{len(image_files)} 张图片")

        # 生成汇总报告

        if results:

            summary_file = os.path.join("final_results", "processing_summary.txt")

            with open(summary_file, 'w', encoding='utf-8') as f:

                f.write("图片处理汇总报告\n")

                f.write("=" * 50 + "\n\n")

                f.write(f"总图片数: {len(image_files)}\n")

                f.write(f"成功处理: {len(results)}\n\n")

                f.write("处理结果文件:\n")

                for result in results:
                    f.write(f"- {os.path.basename(result['saved_file'])}\n")

            print(f"📋 汇总报告: {summary_file}")

    else:
        print("无效选择!")


if __name__ == "__main__":
    main()